{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c9d89b5c",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70996d8a",
   "metadata": {},
   "source": [
    "# IBM watsonx.ai\n",
    "\n",
    ">WatsonxLLM is a wrapper for IBM [watsonx.ai](https://www.ibm.com/products/watsonx-ai) foundation models.\n",
    "\n",
    "The aim of these examples is to show how to communicate with `watsonx.ai` models using the `LlamaIndex` LLMs API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea35b2b7",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "Install the `llama-index-llms-ibm` package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1fff4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU llama-index-llms-ibm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f406e092",
   "metadata": {},
   "source": [
    "The cell below defines the credentials required to work with watsonx Foundation Model inferencing.\n",
    "\n",
    "**Action:** Provide the IBM Cloud user API key. For details, see\n",
    "[Managing user API keys](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d572a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "watsonx_api_key = getpass()\n",
    "os.environ[\"WATSONX_APIKEY\"] = watsonx_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c59782a7",
   "metadata": {},
   "source": [
    "Additionally, you can pass additional secrets as an environment variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98c573c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"WATSONX_URL\"] = \"your service instance url\"\n",
    "os.environ[\"WATSONX_TOKEN\"] = \"your token for accessing the CPD cluster\"\n",
    "os.environ[\"WATSONX_PASSWORD\"] = \"your password for accessing the CPD cluster\"\n",
    "os.environ[\"WATSONX_USERNAME\"] = \"your username for accessing the CPD cluster\"\n",
    "os.environ[\n",
    "    \"WATSONX_INSTANCE_ID\"\n",
    "] = \"your instance_id for accessing the CPD cluster\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36acbef",
   "metadata": {},
   "source": [
    "## Load the model\n",
    "\n",
    "You might need to adjust model `parameters` for different models or tasks. For details, refer to [Available MetaNames](https://ibm.github.io/watsonx-ai-python-sdk/fm_model.html#metanames.GenTextParamsMetaNames)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "407cd500",
   "metadata": {},
   "outputs": [],
   "source": [
    "temperature = 0.5\n",
    "max_new_tokens = 50\n",
    "additional_params = {\n",
    "    \"decoding_method\": \"sample\",\n",
    "    \"min_new_tokens\": 1,\n",
    "    \"top_k\": 50,\n",
    "    \"top_p\": 1,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b586538",
   "metadata": {},
   "source": [
    "Initialize the `WatsonxLLM` class with the previously set parameters.\n",
    "\n",
    "\n",
    "**Note**: \n",
    "\n",
    "- To provide context for the API call, you must pass the `project_id` or `space_id`. To get your project or space ID, open your project or space, go to the **Manage** tab, and click **General**. For more information see: [Project documentation](https://www.ibm.com/docs/en/watsonx-as-a-service?topic=projects) or [Deployment space documentation](https://www.ibm.com/docs/en/watsonx/saas?topic=spaces-creating-deployment).\n",
    "- Depending on the region of your provisioned service instance, use one of the urls listed in [watsonx.ai API Authentication](https://ibm.github.io/watsonx-ai-python-sdk/setup_cloud.html#authentication).\n",
    "\n",
    "In this example, we’ll use the `project_id` and Dallas URL.\n",
    "\n",
    "\n",
    "You need to specify the `model_id` that will be used for inferencing. You can find the list of all the available models in [Supported foundation models](https://ibm.github.io/watsonx-ai-python-sdk/fm_model.html#ibm_watsonx_ai.foundation_models.utils.enums.ModelTypes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359898de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ibm import WatsonxLLM\n",
    "\n",
    "watsonx_llm = WatsonxLLM(\n",
    "    model_id=\"ibm/granite-13b-instruct-v2\",\n",
    "    url=\"https://us-south.ml.cloud.ibm.com\",\n",
    "    project_id=\"PASTE YOUR PROJECT_ID HERE\",\n",
    "    temperature=temperature,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    additional_params=additional_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2202f4e0",
   "metadata": {},
   "source": [
    "Alternatively, you can use Cloud Pak for Data credentials. For details, see [watsonx.ai software setup](https://ibm.github.io/watsonx-ai-python-sdk/setup_cpd.html).    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "243ecccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "watsonx_llm = WatsonxLLM(\n",
    "    model_id=\"ibm/granite-13b-instruct-v2\",\n",
    "    url=\"PASTE YOUR URL HERE\",\n",
    "    username=\"PASTE YOUR USERNAME HERE\",\n",
    "    password=\"PASTE YOUR PASSWORD HERE\",\n",
    "    instance_id=\"openshift\",\n",
    "    version=\"4.8\",\n",
    "    project_id=\"PASTE YOUR PROJECT_ID HERE\",\n",
    "    temperature=temperature,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    additional_params=additional_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96ed13d4",
   "metadata": {},
   "source": [
    "Instead of `model_id`, you can also pass the `deployment_id` of the previously tuned model. The entire model tuning workflow is described in [Working with TuneExperiment and PromptTuner](https://ibm.github.io/watsonx-ai-python-sdk/pt_working_with_class_and_prompt_tuner.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08e66c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "watsonx_llm = WatsonxLLM(\n",
    "    deployment_id=\"PASTE YOUR DEPLOYMENT_ID HERE\",\n",
    "    url=\"https://us-south.ml.cloud.ibm.com\",\n",
    "    project_id=\"PASTE YOUR PROJECT_ID HERE\",\n",
    "    temperature=temperature,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    additional_params=additional_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25ecbd1",
   "metadata": {},
   "source": [
    "## Create a Completion\n",
    "Call the model directly using a string type prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc076c56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A generative AI is a computer program that can create new text, images, or other types of content. These programs are trained on large datasets of existing content, and they use that data to generate new content that is similar to the training data.\n"
     ]
    }
   ],
   "source": [
    "response = watsonx_llm.complete(\"What is a Generative AI?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf3098f",
   "metadata": {},
   "source": [
    "From the `CompletionResponse`, you can also retrieve a raw response returned by the service:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9414bc2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'model_id': 'ibm/granite-13b-instruct-v2', 'created_at': '2024-05-20T07:11:57.984Z', 'results': [{'generated_text': 'A generative AI is a computer program that can create new text, images, or other types of content. These programs are trained on large datasets of existing content, and they use that data to generate new content that is similar to the training data.', 'generated_token_count': 50, 'input_token_count': 7, 'stop_reason': 'max_tokens', 'seed': 494448017}]}\n"
     ]
    }
   ],
   "source": [
    "print(response.raw)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e79dc3",
   "metadata": {},
   "source": [
    "You can also call a model that provides a prompt template:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb291e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A loan is a sum of money that is borrowed to buy something, such as a house or a car. The borrower must repay the loan plus interest. The interest is a fee charged for using the money. The interest rate is the amount of\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core import PromptTemplate\n",
    "\n",
    "template = \"What is {object} and how does it work?\"\n",
    "prompt_template = PromptTemplate(template=template)\n",
    "\n",
    "prompt = prompt_template.format(object=\"a loan\")\n",
    "\n",
    "response = watsonx_llm.complete(prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da1eb2b1",
   "metadata": {},
   "source": [
    "## Calling `chat` with a list of messages\n",
    "Create `chat` completions by providing a list of messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd9980c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: I am an AI assistant.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are an AI assistant\"),\n",
    "    ChatMessage(role=\"user\", content=\"Who are you?\"),\n",
    "]\n",
    "response = watsonx_llm.chat(\n",
    "    messages, max_new_tokens=20, decoding_method=\"greedy\"\n",
    ")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50083c7f",
   "metadata": {},
   "source": [
    "Note that we changed the `max_new_tokens` parameter to `20` and the `decoding_method` parameter to `greedy`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c9da33",
   "metadata": {},
   "source": [
    "## Streaming the model output \n",
    "\n",
    "Stream the model's response:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f63166a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I like New York because it is the city of dreams. You can achieve anything you want here. "
     ]
    }
   ],
   "source": [
    "for chunk in watsonx_llm.stream_complete(\n",
    "    \"Describe your favorite city and why it is your favorite.\"\n",
    "):\n",
    "    print(chunk.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "373db553",
   "metadata": {},
   "source": [
    "Similarly, to stream the `chat` completions, use the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "308e5ba3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I am an AI assistant."
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are an AI assistant\"),\n",
    "    ChatMessage(role=\"user\", content=\"Who are you?\"),\n",
    "]\n",
    "\n",
    "for chunk in watsonx_llm.stream_chat(messages):\n",
    "    print(chunk.delta, end=\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
